from celery import shared_task
from copy import deepcopy
from .app import app
from .planner import Planner


def fake_compute(payload):
    return deepcopy(payload)


def real_compute(payload):
    data = deepcopy(payload)

    planner = Planner(
        algo_type=data['config']['algo_type'],
        num_beams=data['config']['num_beams'],
        selected_container=data['config']['container'],
    )

    return planner.compute(data['input_dataset'])


@shared_task(name='zmm.do_compute')
def on_task_compute(payload):
    try:
        app.send_task('zmm.do_start', [payload])

        # result = fake_compute(payload)
        result = real_compute(payload)

        payload['result_dataset'] = result
        app.send_task('zmm.do_finish', [payload])

    except Exception as e:
        app.send_task('zmm.do_fail', [payload])
        raise e
